"""
该文件仅用于定义LSTM模型
"""
import tensorflow as tf
import config


N_HIDDEN_UNITS = config.N_HIDDEN_UNITS
BATCH_SIZE = config.LSTM_BATCH_SIZE


# 取最后一个值
def _last_relevant(output, length):
    batch_size = tf.shape(output)[0]
    max_length = tf.shape(output)[1]
    out_size = int(output.get_shape()[2])
    index = tf.range(0, batch_size) * max_length + (length - 1)
    flat = tf.reshape(output, [-1, out_size])
    result = tf.gather(flat, index)
    return result


# 模型
def GRU(X, weights, biases, true_length):
    with tf.variable_scope('init_name', initializer=tf.orthogonal_initializer()):   # 正交初始化
        cell = tf.contrib.rnn.GRUCell(N_HIDDEN_UNITS)   # 定义神经元节点
        # tf.constant_initializer(0.0)
        init_state = tf.get_variable('init_state', [1, N_HIDDEN_UNITS], initializer=tf.constant_initializer(0.0))
        init_state = tf.tile(init_state, [BATCH_SIZE, 1])
        outputs, states = tf.nn.dynamic_rnn(
                cell, X, dtype=tf.float32, sequence_length=true_length, initial_state=init_state)
        outputs = tf.nn.dropout(outputs, 0.5)     # dropout层，防止或减轻过拟合
        last = _last_relevant(outputs, true_length)
        results = tf.matmul(last, weights['out']) + biases['out']
        return results
